import os
import shutil
import glob

rootDir = '/data/ICE/RGBCollection'
rgb_imageDir = os.path.join(rootDir, 'rgb_images')

base_rootDir = os.path.join(rootDir, 'base')
cam0_rootDir = os.path.join(rootDir, 'cam0')
cam1_rootDir = os.path.join(rootDir, 'cam1')
base_imageDir = os.path.join(base_rootDir, 'images')
base_xmlDir = os.path.join(base_rootDir, 'xml')
cam0_imageDir = os.path.join(cam0_rootDir, 'images')
cam0_xmlDir = os.path.join(cam0_rootDir, 'xml')
cam1_imageDir = os.path.join(cam1_rootDir, 'images')
cam1_xmlDir = os.path.join(cam1_rootDir, 'xml')

extract_path = '/data/ICE/DataSetCollection/batch2_batch3_tiny'
extract_images_path = os.path.join(extract_path, 'JPEGImages')
extract_xml_path = os.path.join(extract_path, 'xml')
extract_images_name = os.listdir(extract_images_path)
rgb_images_name = os.listdir(rgb_imageDir)
baseimages = []
basexmls = []
cam0_images = []
cam0_xmls = []
cam1_images = []
cam1_xmls = []

for rgb_name in rgb_images_name:
    basename = rgb_name
    cam0_basename = "cam0" + rgb_name
    cam1_basename = "cam1" + rgb_name
    cam0_basename_1 = "cam0_" + rgb_name
    cam1_basename_1 = "cam1_" + rgb_name
    if basename in extract_images_name:
        baseimages.append(basename)
    elif cam0_basename in extract_images_name:
        cam0_images.append(cam0_basename)
    elif cam0_basename_1 in extract_images_name:
        cam0_images.append(cam0_basename_1)
    elif cam1_basename in extract_images_name:
        cam1_images.append(cam1_basename)
    elif cam1_basename_1 in extract_images_name:
        cam1_images.append(cam1_basename_1)

for baseimage in baseimages:
    src_image = os.path.join(extract_images_path, baseimage)
    src_xml = os.path.join(extract_xml_path, baseimage[:-4] + '.xml')
    shutil.copy(src_image, base_imageDir)
    shutil.copy(src_xml, base_xmlDir)
for cam0_image in cam0_images:
    src_image = os.path.join(extract_images_path, cam0_image)
    src_xml = os.path.join(extract_xml_path, cam0_image[:-4] + '.xml')
    shutil.copy(src_image, cam0_imageDir)
    shutil.copy(src_xml, cam0_xmlDir)
for cam1_image in cam1_images:
    src_image = os.path.join(extract_images_path, cam1_image)
    src_xml = os.path.join(extract_xml_path, cam1_image[:-4] + '.xml')
    shutil.copy(src_image, cam1_imageDir)
    shutil.copy(src_xml, cam1_xmlDir)

